import gymnasium as gym

# 创建环境
env = gym.make('CartPole-v1')
env = gym.make("CartPole-v1", render_mode="human")
# PID参数优化（新增积分项并调整权重）
Kp_theta = 20.0    # 角度比例系数（保持杆直立）
Kd_theta = 3.0     # 角度微分系数（抑制振荡）
Kp_x = -2.0        # 位置比例系数（更强的反向控制）
Kd_x = -0.5        # 位置微分系数（更快的速度纠正）
Ki_x = -0.02       # 位置积分系数（消除稳态误差）


# 重置环境并初始化状态
state, info = env.reset()
done = False
total_reward = 0
integral_x = 0.0
# 运行控制循环
while not done:
    x, x_dot, theta, theta_dot = state
    
     # 计算位置积分项（累积误差）
    integral_x += x
    
    # PID控制公式（含积分项）
    control = (
        Kp_theta * theta + Kd_theta * theta_dot +  # 角度控制
        Kp_x * x + Kd_x * x_dot + Ki_x * integral_x  # 位置控制
    )
    
    # 选择动作：control > 0时向右，否则向左
    action = 1 if control > 0 else 0
    
    # 执行动作并获取新状态

    state, reward, done, _, _ = env.step(action)
    total_reward += reward
    if (int(total_reward) % 100 == 0) :
        print(state)
    # 渲染环境（可选）
    env.render()

# 关闭环境
env.close()
print(f"Total Reward: {total_reward}")